import librosa
import numpy as np
import matplotlib.pyplot as plt
import warnings
import soundfile as sf
import math

class Augmentation():
    def __init__(self,sample_rate):
        self.sample_rate = sample_rate

    def _read_wav(self,wave_file):
        data,_ = librosa.load(wave_file,self.sample_rate)
        channel = len(data)
        if channel ==2:
            data = np.mean(data,axis=1,dtype=data.dtype)
        return data
    def _save_wav(self,save_file,data):
        return sf.write(save_file,data.astype(np.float32),16000,subtype="PCM_24")

    def plot_wav(self,data):
        plt.figure()
        plt.plot(data)
        plt.show()

    def snr2noise(self,clean_wav, noise, SNR):
        """
        :param clean_wav: 纯净语音
        :param noise: 噪音
        :param SER: 指定的SNR
        :return: 根据指定的SNR求带噪语音(纯净语音+噪声)
        """
        if len(noise) >= len(clean_wav):
            noise = noise[:len(clean_wav)]

        else:
            mul = math.ceil(len(clean_wav)/len(noise))
            noise = noise*mul
            noise = noise[:len(clean_wav)]


        p_clean = np.mean(clean_wav ** 2)  # 纯净语音功率
        p_noise = np.mean(noise ** 2)  # 噪声功率

        scalar = np.sqrt(p_clean / (10 ** (SNR / 10)) / (p_noise + np.finfo(np.float32).eps))
        noisy = clean_wav + scalar * noise

        return noisy

    def add_noise(self,clean_wav, noise,w=0.004):
        # w：噪声因子
        if len(noise) >= len(clean_wav):
            noise = noise[:len(clean_wav)]

        else:
            mul = math.ceil(len(clean_wav)/len(noise))
            noise = noise*mul
            noise = noise[:len(clean_wav)]
        output = clean_wav + w * noise
        return output

    def dB_gain(self,clean_wav, dB):
        """
        :param wav: 语音
        :param dB: 音量
        :return:返回以指定dB增益后的语音
        """
        wav_p = np.mean(clean_wav ** 2)  # 平均功率
        scalar = np.sqrt(10 ** (dB / 10) / (wav_p + np.finfo(np.float32).eps))
        clean_wav *= scalar
        return clean_wav

if __name__=="__main__":
    filename = "/Users/qianghaozhang/Desktop/科沃斯/智能语音数据/1.5.1.wav"
    target_file = "/Users/qianghaozhang/Desktop/科沃斯/智能语音数据/1.5.1_t.wav"
    noice_file ="/Users/qianghaozhang/Desktop/科沃斯/智能语音数据/_background_noise_/doing_the_dishes.wav"
    aug = Augmentation(16000)
    data = aug._read_wav(filename)
    noice = aug._read_wav(noice_file)
    mix_data = aug.dB_gain(data,2)
    aug._save_wav(target_file,mix_data)

